import torch
import torch.nn as nn

# 传统线性模型（如逻辑回归回归）无法捕捉特征间的交互关系，而多项式模型（如二阶交叉项）在数据稀疏时会面临两个问题：
# 特征组合（如 “性别 = 女”×“年龄段 = 20-30”）在训练数据中可能极少出现，导致参数无法学习
# 特征组合数量随特征数呈平方增长（O(n2)），计算成本极高

class DeepFM(nn.Module):
    """DeepFM模型实现"""

    def __init__(self, cat_dims, num_cols_count, embedding_dim=8, hidden_dims=None):
        super(DeepFM, self).__init__()
        if hidden_dims is None:
            hidden_dims = [512, 256, 128, 64]
        self.cat_dims = cat_dims
        # 为什么可以设置不同的嵌入维度？
        # DeepFM 的核心是 “factorization machine（FM）部分 + deep neural network（DNN）部分”，两者共享嵌入层的输出：
        # FM 部分：需要将不同类别特征的嵌入向量做 pairwise 交互（点积运算），只要嵌入向量的维度一致即可参与运算（与单个特征的嵌入维度无关，只要求所有特征的嵌入维度相同）。
        # DNN 部分：需要将所有嵌入向量拼接后输入网络，同样只要求所有嵌入向量的维度一致（否则无法拼接）。
        # 这里的关键是：同一模型中所有类别特征的嵌入向量必须具有相同的维度，但这个维度可以针对不同特征 “动态指定”，只要最终保持一致即可。
        self.embedding_dim = embedding_dim

        # 嵌入层（Embedding Layer） 是一种将离散的类别型特征（如用户 ID、物品 ID、性别等）转换为连续的低维向量表示的技术。它解决了类别特征直接使用 one-hot
        # 编码导致的维度爆炸增和稀疏性问题，同时能捕捉类别间的潜在关系
        # nn.Embedding(dim, embedding_dim)：PyTorch 的嵌入层类
        # dim：该类别特征的总类别数（如性别有 2 类，dim=2）
        # embedding_dim：嵌入向量的维度（如 8、16，超参数，需手动设置）
        # 内部维护一个形状为 [dim, embedding_dim] 的可训练矩阵，每个类别对应矩阵中的一行
        # nn.ModuleList：管理多个嵌入层的容器
        # 因为不同类别特征（如用户 ID、物品 ID）的类别数不同，需要为每个特征创建独立的嵌入层
        # 例如 cat_dims = [1000, 500, 2] 表示 3 个类别特征，分别有 1000、500、2 个类别，代码会创建 3 个嵌入层
        # 稀疏矩阵降维
        self.embeddings = nn.ModuleList([
            nn.Embedding(dim, embedding_dim) for dim in cat_dims
        ])

        # FM线性层
        self.fm_linear = nn.Linear(sum(cat_dims) + num_cols_count, 1)

        # DNN部分
        input_dnn_dim = embedding_dim * len(cat_dims) + num_cols_count
        self.dnn = nn.Sequential()
        for i, hidden_dim in enumerate(hidden_dims):
            # 定义全连接层
            self.dnn.add_module(f'linear_{i}', nn.Linear(input_dnn_dim, hidden_dim))
            # 激活函数
            self.dnn.add_module(f'relu_{i}', nn.ReLU())
            # 丢弃特征 防止过拟合
            self.dnn.add_module(f'dropout_{i}', nn.Dropout(0.3))
            input_dnn_dim = hidden_dim

        # 输出层
        self.output = nn.Linear(hidden_dims[-1] + 1, 1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, cat_features, num_features):
        # 类别特征嵌入向量
        embeddings = [emb(cat_features[:, i]) for i, emb in enumerate(self.embeddings)]
        # 保留原始嵌入列表用于FM交互项计算（每个元素形状: [batch_size, embedding_dim]）

        # 拼接嵌入向量用于DNN（形状: [batch_size, num_cat_features * embedding_dim]）
        cat_emb = torch.cat(embeddings, dim=1)

        # 拼接所有特征（用于DNN）
        all_features = torch.cat([cat_emb, num_features], dim=1)

        # FM线性项
        one_hot_cat = torch.cat([
            nn.functional.one_hot(cat_features[:, i], num_classes=dim).float()
            for i, dim in enumerate(self.cat_dims)
        ], dim=1)
        fm_linear_input = torch.cat([one_hot_cat, num_features], dim=1)
        fm_linear_part = self.fm_linear(fm_linear_input)  # 形状: [batch_size, 1]

        # FM交互项
        # 将嵌入列表堆叠为三维张量 [batch_size, num_cat_features, embedding_dim]
        emb_stack = torch.stack(embeddings, dim=1)  # 关键修复：保持特征维度

        # 计算求和的平方
        sum_emb = torch.sum(emb_stack, dim=1)  # [batch_size, embedding_dim]
        sum_square = sum_emb ** 2  # [batch_size, embedding_dim]

        # 计算平方的求和
        square_emb = emb_stack ** 2  # [batch_size, num_cat_features, embedding_dim]
        square_sum = torch.sum(square_emb, dim=1)  # [batch_size, embedding_dim]

        # 计算交互项（此时dim=1有效）
        fm_interaction_part = 0.5 * torch.sum(sum_square - square_sum, dim=1, keepdim=True)  # [batch_size, 1]
        fm_output = fm_linear_part + fm_interaction_part  # [batch_size, 1]

        # DNN输出
        dnn_output = self.dnn(all_features)  # 形状: [batch_size, dnn_output_dim]

        # 最终输出
        combined = torch.cat([fm_output, dnn_output], dim=1)  # 拼接FM和DNN输出
        output = self.output(combined)  # 压缩到1维
        return self.sigmoid(output)  # 二分类概率
